#include <math.h>
#include <map>
#include <limits>
#include <vector>

#include "sensor_model.h"

//#define DEBUG_FLAG
#define FILL_VALUE 0.0

#define SCALE 0.01

#define BEAM_MODEL 0
#define LF_MODEL   1

SensorModel::SensorModel(ros::NodeHandle& nh, ros::NodeHandle& ph, boost::mt19937& g)
:n(nh), pn(ph),ptree(NULL),
 gen(g), uniform_2pi(g, boost::uniform_real<>(-M_PI, M_PI))
{
	ros::NodeHandle pn("~");

	pn.param("model_type",   model_type,	BEAM_MODEL);
	pn.param("z_hit",    	 z_hit, 	 	0.5);
	pn.param("z_short", 	 z_short,		0.05);
	pn.param("z_max",		 z_max,			0.05);
	pn.param("z_rand",   	 z_rand,     	0.5);
	pn.param("sigma_hit", 	 sigma_hit, 	0.2);
	pn.param("lambda_short", lambda_short,  0.1);
	pn.param("max_occ_dist", max_occ_dist, 	2.0);
	pn.param("laserXoffset", x_offset,		0.25);
	pn.param("laserYoffset", y_offset,		0.0);
	pn.param("raycastStep",  raycastStep,   0.1);
	pn.param("sweep_step",   sweepStep,     5);



	// This is where we get the map
	map_client = n.serviceClient<nav_msgs::GetMap>("static_map");
	requestMap();

	pcl_pub = n.advertise<geometry_msgs::PoseArray>("beam", 2);

	if (model_type ==LF_MODEL)
	{
		ROS_INFO("Using Likelihood Model");
		ROS_INFO("z_hit %.3lf z_rand %.3lf sigma_hit %.3lf maxDist %.3lf",
				  z_hit, z_rand, sigma_hit, max_occ_dist);
		// build likelyhood_field map
		ptree = kd_create( 2 );

		ROS_INFO("Initializing likelihood map");
		initializeMap();

		ROS_INFO("Likelihood map generated.");
	}
	else
	{
		ROS_INFO("Using Beam Model");
		ROS_INFO("z_hit %.3lf z_short %.3lf z_rand %.3lf z_max %.3lf sigma_hit %.3lf",
				  z_hit, z_short, z_rand, z_max, sigma_hit);
		initializeMap();
		ROS_INFO("IM fine");
	}

}

SensorModel::~SensorModel()
{
	if (model_type ==LF_MODEL)
		kd_free( ptree );
}

void SensorModel::publish_bestMeasurement(geometry_msgs::Pose& pose, std::vector<int>& data)
{
	geometry_msgs::PoseArray beam;
	beam.header.frame_id = "map";
	beam.header.stamp = ros::Time::now();
	geometry_msgs::Pose p;
	double yaw = tf::getYaw(pose.orientation);
	double x_sense = pose.position.x + x_offset* cos(yaw) - y_offset* sin(yaw);
	double y_sense = pose.position.y + y_offset* cos(yaw) + x_offset* sin(yaw);

	for (unsigned int i=0; i< data.size(); ++i)
	{
		double range = SCALE*data[i];
		if ( range >= max_occ_dist)
		{
			continue;
		}

		double x_query = x_sense+range*cos(yaw+ i*M_PI/180.0 -M_PI*0.5);// convert cm to m;
		double y_query = y_sense+range*sin(yaw+ i*M_PI/180.0 -M_PI*0.5);

		p.position.x = x_query;
		p.position.y = y_query;
//		p.position.z = dist;//10*q;
		p.orientation = pose.orientation;
		beam.poses.push_back(p);
	}
	pcl_pub.publish(beam);
}

double SensorModel::getWeight(geometry_msgs::Pose& pose, std::vector<int>& data)
{
	if (model_type == BEAM_MODEL)
		return getWeightBeam(pose, data);
	else
		return getWeightLF(pose, data);
}

double SensorModel::getWeightBeam(geometry_msgs::Pose& pose, std::vector<int>& data)
{
	double yaw = tf::getYaw(pose.orientation);
	double weight = 1.0;
	double x_sense = pose.position.x + x_offset* cos(yaw) - y_offset* sin(yaw);
	double y_sense = pose.position.y + y_offset* cos(yaw) + x_offset* sin(yaw);

	for (int i=0; i< data.size(); i+=sweepStep)
	{
		double range = SCALE*data[i];

		double x_end = x_sense+ max_occ_dist*cos(yaw+ i*M_PI/180.0 -M_PI*0.5);
		double y_end = y_sense+ max_occ_dist*sin(yaw+ i*M_PI/180.0 -M_PI*0.5);

		double zk = raytrace(x_sense/meta_data.resolution, y_sense/meta_data.resolution,
							 x_end/meta_data.resolution,   y_end/meta_data.resolution);

		double z = range -zk;
		double phit = z_hit* exp(-(z * z) / (2 * sigma_hit * sigma_hit));

		double pshort = 0.0;
		if (z < 0.0)
			pshort = z_short * lambda_short * exp(-lambda_short*range);

		double pmax = 0.0;
		if (range == max_occ_dist)
			pmax = z_max;

		double pz = phit + pshort + pmax + z_rand / max_occ_dist;

//		ROS_INFO("zk%.8lf vs range %.6lf pz %.8lf phit %.8lf pshort %.8lf", zk, range, pz,phit,pshort);
		weight *= pow(pz, 0.1/sweepStep);
//		weight *= (1.0 + pz*pz*pz);
	}
	return weight;
}

double SensorModel::getWeightLF(geometry_msgs::Pose& pose, std::vector<int>& data)
{
	double yaw = tf::getYaw(pose.orientation);
	double weight = 1.0;
	double x_sense = pose.position.x + x_offset* cos(yaw) - y_offset* sin(yaw);
	double y_sense = pose.position.y + y_offset* cos(yaw) + x_offset* sin(yaw);

	for (unsigned int i=0; i< data.size(); ++i)
	{
		double range = SCALE*data[i];
		double dist;
		if ( range >= max_occ_dist)
		{
			continue;
		}

		double x_query = x_sense+range*cos(yaw+ i*M_PI/180.0 -M_PI*0.5);// convert cm to m;
		double y_query = y_sense+range*sin(yaw+ i*M_PI/180.0 -M_PI*0.5);

		if ( queryPtOnMap(x_query, y_query) )
		{
			dist = LF_map[xy2Idx(x_query, y_query)];
		}
		else
		{
			dist = max_occ_dist;//minDist(x_query, y_query);
		}

		double q = (z_hit*prob(dist, sigma_hit) + z_rand/max_occ_dist);

		q= pow(q, 0.25);

		assert(q <= 1.0);
		assert(q >= 0.0);
//		ROS_INFO("x: %.4lf, y: %.4lf, dist: %.4lf, q: %.6lf", x_query, y_query, dist, q);
		weight *=q;
	}
	return weight;
}

void SensorModel::initializeMap()
{
	// Initialize map dimension
	X_max = meta_data.width * meta_data.resolution;
	Y_max = meta_data.height* meta_data.resolution;

	std::vector<int > query;

	size_t size = Omap.data.size();

	// Step 1: build kdtree and find where we need to make the map
	for (size_t i=0; i < size; ++i)
	{
		switch (Omap.data[i])
		{
		case 0: // empty cells
//		case -2:// border cells
			query.push_back(i);
			break;
		case 100:// concrete obstacles
			if (model_type ==LF_MODEL)
			{
				double obstacle[2];
				obstacle[0] = getXmapFrame(i);
				obstacle[1] = getYmapFrame(i);
				kd_insert(ptree, obstacle, NULL);
			}
			break;
		case -1:// unknown area
		default:
			break;
		}
	}

	if (model_type ==LF_MODEL)
	{
		// Step 2: build the likelihood field map
		LF_map.resize(meta_data.width*meta_data.height);
		for(size_t i=0; i< LF_map.size(); ++i)
			LF_map[i] = FILL_VALUE;
	}

	double x, y;
	for(size_t j=0; j< query.size(); ++j)
	{
		x = getXmapFrame(query[j]);
		y = getYmapFrame(query[j]);
		free_area.push_back(std::pair<double, double>(x, y));
		if (model_type ==LF_MODEL)
		{
			LF_map[query[j]] = prob(minDist(x, y), sigma_hit);
		}
	}

#ifdef DEBUG_FLAG
	if (model_type ==LF_MODEL)
	{
		FILE* debug_file= fopen("lf_map.dat", "w");
		if (debug_file ==0)
			ROS_ERROR("Failed to open file lf_map.dat");

		for(int i=0; i< 800; i++)
		{
			for(int j=0; j< 800; j++)
				fprintf(debug_file, "%.4lf ", LF_map[i*800+j]);
			fprintf(debug_file, "\n");
		}
		fclose(debug_file);
		ROS_INFO("Write to file");
	}
#endif
}

void SensorModel::addNRandParticles(const int N, std::vector<geometry_msgs::Pose>& cloud)
{

//	static boost::uniform_real<> uniform_2pi(-M_PI, M_PI);

	ROS_WARN("Generate %d particles", N);

	for(int i=0; i< N; ++i)
	{
		geometry_msgs::Pose p = getOneRandParticle();
		cloud.push_back(p);
	}
}

geometry_msgs::Pose SensorModel::getOneRandParticle()
{
	static boost::uniform_int<>  uniform_Int(0, free_area.size());
	geometry_msgs::Pose p;
	int j = uniform_Int(gen);
	p.position.x = free_area[j].first;
	p.position.y = free_area[j].second;
	double angle = uniform_2pi();
	assert(fabs(angle) < (2*M_PI)+1e-4);
	p.orientation= tf::createQuaternionMsgFromYaw(angle);
	return p;
}

bool SensorModel::queryPtOnMap(double x, double y)
{
	return (x>=0 && x<= X_max && y>=0 && y<= Y_max);
}

int SensorModel::xy2Idx(double x, double y)
{
	return static_cast<int>(x/meta_data.resolution)+static_cast<int>(y/meta_data.resolution)*meta_data.width;
}

double SensorModel::minDist(double x, double y)
{
	struct kdres *presults;
	double dist;
	double pt[2] = {x, y};
	double pos[2];

	presults = kd_nearest( ptree, pt );
	kd_res_item( presults, pos );

	if ( presults == 0)
	{
		ROS_ERROR("Couldn't find a nearest neighbor");
		dist = 50;
	}
	else
		dist = sqrt( (pt[0]-pos[0])*(pt[0]-pos[0])+(pt[1]-pos[1])*(pt[1]-pos[1]));

	kd_res_free( presults );
	return dist;
}

double SensorModel::prob(double d, double sigma)
{
	return (exp(-0.5*(d*d)/(sigma*sigma)) / (sigma* sqrt(2*M_PI)));
}

double SensorModel::getYmapFrame(int grid_y)
{
	return meta_data.resolution* (grid_y/meta_data.width);
}

double SensorModel::getXmapFrame(int grid_x)
{
	return meta_data.resolution* (grid_x%meta_data.width);
}

void SensorModel::requestMap()
{
	nav_msgs::GetMap map_srv;

	int count = 5;
	while(!map_client.call(map_srv))
	{
		ROS_ERROR("Call to map_server failed");
		sleep(1);
		if(--count ==0)
			exit(1);
	}
	Omap = map_srv.response.map;
	meta_data = Omap.info;

}

double SensorModel::raytrace(double x0, double y0, double x1, double y1)
{
//	ROS_INFO("Start at %lf %lf, ends at %lf %lf", x0, y0, x1, y1);
	// x0, x1, y0, y1 are all in map frame now
    double dx = fabs(x1 - x0);
    double dy = fabs(y1 - y0);

    int x = int(floor(x0));
    int y = int(floor(y0));

    int n = 1;
    int x_inc, y_inc;
    double error;

    if (dx == 0)
    {
        x_inc = 0;
        error = std::numeric_limits<double>::infinity();
    }
    else if (x1 > x0)
    {
        x_inc = 1;
        n += int(floor(x1)) - x;
        error = (floor(x0) + 1 - x0) * dy;
    }
    else
    {
        x_inc = -1;
        n += x - int(floor(x1));
        error = (x0 - floor(x0)) * dy;
    }

    if (dy == 0)
    {
        y_inc = 0;
        error -= std::numeric_limits<double>::infinity();
    }
    else if (y1 > y0)
    {
        y_inc = 1;
        n += int(floor(y1)) - y;
        error -= (floor(y0) + 1 - y0) * dx;
    }
    else
    {
        y_inc = -1;
        n += y - int(floor(y1));
        error -= (y0 - floor(y0)) * dx;
    }

    for (; n > 0; --n)
    {
    	if(x<0 || x > meta_data.width || y<0 || y>meta_data.height)
    	{
//    		ROS_ERROR("at x %d y %d", x, y);
    		return max_occ_dist;//meta_data.resolution*sqrt((x-x0)*(x-x0)+(y-y0)*(y-y0));
    	}

    	else if(Omap.data[x+meta_data.width*y] == 100)
        {
//        	ROS_ERROR("at x %d y %d O =%d", x, y, Omap.data[x+meta_data.width*y]);
        	// return a distance
        	return meta_data.resolution*sqrt((x-x0)*(x-x0)+(y-y0)*(y-y0));
        }

        if (error > 0)
        {
            y += y_inc;
            error -= dx;
        }
        else
        {
            x += x_inc;
            error += dy;
        }
    }
    return max_occ_dist;
}
/* NOT USED in this branch */
/*
 *
double SensorModel::getWeight(geometry_msgs::Pose pose, std::vector<int>& data)
{
	double yaw = tf::getYaw(pose.orientation);
	double weight = 1.0;
	double x_sense = pose.position.x + x_offset* cos(yaw) - y_offset* sin(yaw);
	double y_sense = pose.position.y + y_offset* cos(yaw) + x_offset* sin(yaw);

	for (unsigned int i=0; i< data.size(); ++i)
	{
		double range = SCALE*data[i];
		double dist;
		if ( range >= max_occ_dist)
		{
			continue;
		}

		double x_query = x_sense+range*cos(yaw+ i*M_PI/180.0 -M_PI*0.5);// convert cm to m;
		double y_query = y_sense+range*sin(yaw+ i*M_PI/180.0 -M_PI*0.5);

		if ( queryPtOnMap(x_query, y_query) )
		{
			dist = LF_map[xy2Idx(x_query, y_query)];
		}
		else
		{
			dist = max_occ_dist;//minDist(x_query, y_query);
		}

		double q = (z_hit*prob(dist, sigma_hit) + z_rand/max_occ_dist);

		q= pow(q, 0.25);

		assert(q <= 1.0);
		assert(q >= 0.0);
//		ROS_INFO("x: %.4lf, y: %.4lf, dist: %.4lf, q: %.6lf", x_query, y_query, dist, q);
		weight *=q;
	}
	return weight;
}

double SensorModel::getAllWeight(std::vector<geometry_msgs::Pose>& poses, std::vector<int>& data, std::vector<double>& weights)
{
	double total_weight = 0.0;
	double max_conf = 0.0;
#if SHOW_BEAM
	geometry_msgs::PoseArray beam;
	beam.header.frame_id = "map";
	beam.header.stamp = ros::Time::now();
	geometry_msgs::Pose p;
#endif
	for(int k=0; k < poses.size(); k++)
	{
		geometry_msgs::Pose pose = poses[k];
		double yaw = tf::getYaw(pose.orientation);
		double weight = 1.0;
		double x_sense = pose.position.x + x_offset* cos(yaw) - y_offset* sin(yaw);
		double y_sense = pose.position.y + y_offset* cos(yaw) + x_offset* sin(yaw);

//		ROS_INFO("R@ %.3lf %.3lf BC @ %.3lf %.3lf",  poses[k].position.x, poses[k].position.y, x_sense, y_sense);
		for (unsigned int i=0; i< data.size(); ++i)
		{
			double range = SCALE*data[i];
			double dist;
			if ( range >= max_occ_dist)
			{
				continue;
			}

			double x_query = x_sense+range*cos(yaw+ i*M_PI/180.0 -M_PI*0.5);// convert cm to m;
			double y_query = y_sense+range*sin(yaw+ i*M_PI/180.0 -M_PI*0.5);

#if SHOW_BEAM
	//		geometry_msgs::Pose p;
			p.position.x = x_query;
			p.position.y = y_query;
			p.orientation = pose.orientation;
			beam.poses.push_back(p);
#endif

			if ( queryPtOnMap(x_query, y_query) )
			{
				dist = LF_map[xy2Idx(x_query, y_query)];
			}
			else
			{
				dist = max_occ_dist;//minDist(x_query, y_query);
			}

			double q = (z_hit*prob(dist, sigma_hit) + z_rand/max_occ_dist);

			assert(q <= 1.0);
			assert(q >= 0.0);
//			ROS_INFO("x: %.4lf, y: %.4lf, dist: %.4lf, q: %.6lf", x_query, y_query, dist, q);
			weight *=q;
		}

		if (weight < 1.0 && weight > max_conf)
			max_conf = weight;

		total_weight += weight;
		weights.push_back(weight);
	}

	// normalize the weights;
	for(size_t i=0; i< weights.size(); i++)
		weights[i] /= total_weight;

#if SHOW_BEAM
	pcl_pub.publish(beam);
#endif

	return total_weight;
}
*/
